// Copyright 2022 Chen Jun
// Licensed under the MIT License.

// OpenCV
#include <opencv2/core.hpp>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/opencv.hpp>

// STL
#include <algorithm>
#include <cstddef>
#include <fstream>
#include <map>
#include <string>
#include <vector>

#include "armor_detector/armor.hpp"
#include "armor_detector/number_classifier.hpp"

namespace rm_auto_aim
{
NumberClassifier::NumberClassifier(
  const std::string & model_path, const std::string & label_path, const double thre,
  const std::vector<std::string> & ignore_classes)
: threshold(thre), ignore_classes_(ignore_classes)
{
  net_ = cv::dnn::readNetFromONNX(model_path);

  std::ifstream label_file(label_path);
  std::string line;
  while (std::getline(label_file, line)) {
    class_names_.push_back(line);
  }
}
/*通过透视变换从装甲板对象中提取出数字图像
透视变换的作用：消除装甲板图片的畸变和拉伸，使得装甲板中的数字可以被顺利提取出来，提高准确性可靠性*/
void NumberClassifier::extractNumbers(const cv::Mat & src, std::vector<Armor> & armors)
{
  // Light length in image
  const int light_length = 12;
  // Image size after warp透视变换后的图像大小
  const int warp_height = 28;
  const int small_armor_width = 32;
  const int large_armor_width = 54;
  // Number ROI size 数字感兴趣区域的尺寸
  const cv::Size roi_size(20, 28);

  for (auto & armor : armors) {
    // Warp perspective transform 透视变换
    cv::Point2f lights_vertices[4] = {
      armor.left_light.bottom, armor.left_light.top, armor.right_light.top,
      armor.right_light.bottom};
    // 计算透视变换后的目标顶点坐标
    const int top_light_y = (warp_height - light_length) / 2 - 1;
    const int bottom_light_y = top_light_y + light_length;
    const int warp_width = armor.type == ArmorType::SMALL ? small_armor_width : large_armor_width;
    cv::Point2f target_vertices[4] = {
      cv::Point(0, bottom_light_y),
      cv::Point(0, top_light_y),
      cv::Point(warp_width - 1, top_light_y),
      cv::Point(warp_width - 1, bottom_light_y),
    };
    cv::Mat number_image;
    // 使用 cv::getPerspectiveTransform 函数计算透视变换所需的变换矩阵。
    auto rotation_matrix = cv::getPerspectiveTransform(lights_vertices, target_vertices);
    //使用 cv::warpPerspective 函数将原始图像 src 进行透视变换，获取数字图像。
    cv::warpPerspective(src, number_image, rotation_matrix, cv::Size(warp_width, warp_height));

    // Get ROI  获取感兴趣区域（ROI）——从透视变换后的图像中提取数字区域，得到数字图像
    number_image =
      number_image(cv::Rect(cv::Point((warp_width - roi_size.width) / 2, 0), roi_size));

    // Binarize 二值化处理，以凸显数字的形状。
    cv::cvtColor(number_image, number_image, cv::COLOR_RGB2GRAY);
    cv::threshold(number_image, number_image, 0, 255, cv::THRESH_BINARY | cv::THRESH_OTSU);

    armor.number_img = number_image;
  }
}

/*作用：用于对装甲板中的数字图像进行分类和识别*/
void NumberClassifier::classify(std::vector<Armor> & armors)
{
  for (auto & armor : armors) {
    //克隆装甲板对象中的数字图像。
    cv::Mat image = armor.number_img.clone();

    // Normalize对数字图像进行归一化，将像素值范围缩放到 [0, 1]。
    image = image / 255.0;

    //使用 cv::dnn::blobFromImage 函数将归一化的图像转换为网络模型可接受的输入 blob。
    cv::Mat blob;
    cv::dnn::blobFromImage(image, blob);

    // Set the input blob for the neural network
    net_.setInput(blob);
    // Forward pass the image blob through the model通过网络模型前向传递（forward pass）输入 blob，得到输出结果。
    cv::Mat outputs = net_.forward();//outputs中存储了预测结果

    // 找到输出outputs中概率最大的值
    float max_prob = *std::max_element(outputs.begin<float>(), outputs.end<float>());
    cv::Mat softmax_prob;
    // 对输出结果进行 softmax 运算，得到输出结果的概率分布（不用在意下面这行代码的数学运算）。
    cv::exp(outputs - max_prob, softmax_prob);
    //找到概率分布中概率最高的类别，并获取对应的类别 ID。
    float sum = static_cast<float>(cv::sum(softmax_prob)[0]);
    softmax_prob /= sum;
    
    double confidence;
    cv::Point class_id_point;
    //minMaxLoc 函数查找 softmax_prob.reshape(1, 1)中最大概率值及其位置，概率最高的类别为class_id_point，confidence 变量存储的是该类别对应的概率值
    minMaxLoc(softmax_prob.reshape(1, 1), nullptr, &confidence, nullptr, &class_id_point);
    int label_id = class_id_point.x;
    //更新装甲板对象的数字识别结果和置信度
    armor.confidence = confidence;
    armor.number = class_names_[label_id];
    //将数字识别结果和置信度信息拼接为一个字符串，保存到装甲板对象的 classfication_result 属性中
    std::stringstream result_ss;
    result_ss << armor.number << ": " << std::fixed << std::setprecision(1)
              << armor.confidence * 100.0 << "%";
    armor.classfication_result = result_ss.str();
  }
  //根据置信度和装甲板类型的匹配情况，移除置信度低的装甲板或类型不匹配的装甲板
  armors.erase(
    std::remove_if(
      armors.begin(), armors.end(),
      [this](const Armor & armor) {
        if (armor.confidence < threshold) {
          return true;
        }
        for (const auto & ignore_class : ignore_classes_) //ignore_classes_中存储的是需要排除的负样本的标签（字符串）
        {
          if (armor.number == ignore_class) {
            return true;
          }
        }
        
        bool mismatch_armor_type = false;
        if (armor.type == ArmorType::LARGE) {
          mismatch_armor_type =
            armor.number == "outpost" || armor.number == "2" || armor.number == "guard";
        } else if (armor.type == ArmorType::SMALL) {
          mismatch_armor_type = armor.number == "1" || armor.number == "base";
        }
        return mismatch_armor_type;
      }),
    armors.end());
}

}  // namespace rm_auto_aim
